import numpy as np

def rot_z(deg):
    rad = np.deg2rad(deg)
    c, s = np.cos(rad), np.sin(rad)
    return np.array([[c,-s,0],
                     [s, c,0],
                     [0, 0,1]])

def make_transform_z(translation, yaw_deg):
    tx, ty, tz = translation
    R = rot_z(yaw_deg)
    T = np.eye(4)
    T[:3,:3] = R
    T[:3,3] = [tx, ty, tz]
    return T

def invert_transform(T):
    R = T[:3,:3]; t = T[:3,3]
    T_inv = np.eye(4)
    T_inv[:3,:3] = R.T
    T_inv[:3,3] = -R.T @ t
    return T_inv

def apply_transform(T, points):
    pts = np.atleast_2d(points)
    ones = np.ones((pts.shape[0],1))
    homo = np.hstack((pts, ones))  # (N,4)
    transformed = (T @ homo.T).T
    return transformed[:,:3]

# ------------ 定义默认变换 ------------
T_waist2camera = make_transform_z((47.64571478, 0.0, 462.68178553), -42.0)
T_camera2waist = invert_transform(T_waist2camera)

T_waist2laser = make_transform_z((10.36873357, 0.0, 428.40715287), 2.73146855)
T_laser2waist = invert_transform(T_waist2laser)

# ------------ 坐标轴转换函数 ------------
def laser_axis_transform(point):
    """激光雷达轴转换：绕x轴反转，y和z轴转向
    
    Args:
        point: 一维数组 [x, y, z]，原始激光坐标
        
    Returns:
        一维数组 [x, y, z]，轴转换后的坐标
    """
    x, y, z = point
    # 绕x轴反转：y -> -y, z -> -z
    return np.array([x, -y, -z])

def laser_axis_inverse_transform(point):
    """激光雷达轴逆转换
    
    Args:
        point: 一维数组 [x, y, z]，轴转换后的坐标
        
    Returns:
        一维数组 [x, y, z]，原始激光坐标
    """
    x, y, z = point
    # 逆转换：y -> -y, z -> -z
    return np.array([x, -y, -z])

def camera_axis_transform(point):
    """深度相机轴转换：x是-y轴，y是-z轴，z是x轴
    
    Args:
        point: 一维数组 [x, y, z]，原始相机坐标
        
    Returns:
        一维数组 [x, y, z]，轴转换后的坐标
    """
    x, y, z = point
    # x是-y轴，y是-z轴，z是x轴
    return np.array([z, -x, -y])

def camera_axis_inverse_transform(point):
    """深度相机轴逆转换
    
    Args:
        point: 一维数组 [x, y, z]，轴转换后的坐标
        
    Returns:
        一维数组 [x, y, z]，原始相机坐标
    """
    x, y, z = point
    # 逆转换：从 [z, -x, -y] 恢复到 [x, y, z]
    # 即：x = -y, y = -z, z = x
    return np.array([-y, -z, x])

# ------------ 坐标系转换函数 ------------
def waist_to_camera(point):
    """将腰部坐标系的点转换到相机坐标系
    
    Args:
        point: 一维数组 [x, y, z]，腰部坐标系中的点
        
    Returns:
        一维数组 [x, y, z]，相机坐标系中的点（已经轴转换）
    """
    point_3d = np.array(point).reshape(1, 3)
    # 先进行矩阵转换
    transformed = apply_transform(T_waist2camera, point_3d)
    # 再进行逆向轴转换
    return camera_axis_inverse_transform(transformed[0])

def camera_to_waist(point):
    """将相机坐标系的点转换到腰部坐标系
    
    Args:
        point: 一维数组 [x, y, z]，相机坐标系中的点（已经轴转换）
        
    Returns:
        一维数组 [x, y, z]，腰部坐标系中的点
    """
    # 先进行轴转换
    axis_transformed = camera_axis_transform(point)
    point_3d = np.array(axis_transformed).reshape(1, 3)
    # 再进行矩阵转换
    transformed = apply_transform(T_camera2waist, point_3d)
    return transformed[0]

def waist_to_laser(point):
    """将腰部坐标系的点转换到激光坐标系
    
    Args:
        point: 一维数组 [x, y, z]，腰部坐标系中的点
        
    Returns:
        一维数组 [x, y, z]，激光坐标系中的点（已经轴转换）
    """
    point_3d = np.array(point).reshape(1, 3)
    # 先进行矩阵转换
    transformed = apply_transform(T_waist2laser, point_3d)
    # 再进行逆向轴转换
    return laser_axis_inverse_transform(transformed[0])

def laser_to_waist(point):
    """将激光坐标系的点转换到腰部坐标系
    
    Args:
        point: 一维数组 [x, y, z]，激光坐标系中的点（已经轴转换）
        
    Returns:
        一维数组 [x, y, z]，腰部坐标系中的点
    """
    # 先进行轴转换
    axis_transformed = laser_axis_transform(point)
    point_3d = np.array(axis_transformed).reshape(1, 3)
    # 再进行矩阵转换
    transformed = apply_transform(T_laser2waist, point_3d)
    return transformed[0]

def camera_to_laser(point):
    """将相机坐标系的点转换到激光坐标系
    
    Args:
        point: 一维数组 [x, y, z]，相机坐标系中的点
        
    Returns:
        一维数组 [x, y, z]，激光坐标系中的点
    """
    # 相机 -> 腰部 -> 激光
    waist_point = camera_to_waist(point)
    return waist_to_laser(waist_point)

def laser_to_camera(point):
    """将激光坐标系的点转换到相机坐标系
    
    Args:
        point: 一维数组 [x, y, z]，激光坐标系中的点
        
    Returns:
        一维数组 [x, y, z]，相机坐标系中的点
    """
    # 激光 -> 腰部 -> 相机
    waist_point = laser_to_waist(point)
    return waist_to_camera(waist_point)

# ------------ 验证函数 ------------
def test_axis_transform():
    """测试轴转换函数的正确性"""
    print("测试轴转换函数：")
    
    test_point = [1.0, 2.0, 3.0]
    print(f"原始点: {test_point}")
    
    # 激光轴转换测试
    laser_transformed = laser_axis_transform(test_point)
    laser_recovered = laser_axis_inverse_transform(laser_transformed)
    print(f"激光轴转换: {test_point} -> {laser_transformed} -> {laser_recovered}")
    print(f"激光轴转换误差: {np.linalg.norm(np.array(test_point) - np.array(laser_recovered)):.2e}")
    
    # 相机轴转换测试
    camera_transformed = camera_axis_transform(test_point)
    camera_recovered = camera_axis_inverse_transform(camera_transformed)
    print(f"相机轴转换: {test_point} -> {camera_transformed} -> {camera_recovered}")
    print(f"相机轴转换误差: {np.linalg.norm(np.array(test_point) - np.array(camera_recovered)):.2e}")
    print()

def test_transform():
    """验证坐标系转换函数的正确性"""
    test_points = [
        [0.0, 0.0, 0.0],   # 原点
        [100.0, 0.0, 50.0], # 前上方一点
        [-30.0, 40.0, 10.0] # 任意点
    ]

    print("测试坐标系转换函数（输入输出均为一维数组）\n")
    
    for i, point in enumerate(test_points):
        print(f"测试点 {i+1}: {point}")
        
        # 腰部 -> 相机 -> 腰部往返测试
        camera_point = waist_to_camera(point)
        waist_roundtrip = camera_to_waist(camera_point)
        cam_error = np.linalg.norm(np.array(point) - np.array(waist_roundtrip))
        
        print(f"  腰部 -> 相机: {camera_point}")
        print(f"  相机 -> 腰部: {waist_roundtrip}")
        print(f"  相机往返误差: {cam_error:.2e}")
        
        # 腰部 -> 激光 -> 腰部往返测试
        laser_point = waist_to_laser(point)
        waist_roundtrip2 = laser_to_waist(laser_point)
        laser_error = np.linalg.norm(np.array(point) - np.array(waist_roundtrip2))
        
        print(f"  腰部 -> 激光: {laser_point}")
        print(f"  激光 -> 腰部: {waist_roundtrip2}")
        print(f"  激光往返误差: {laser_error:.2e}")
        
        # 相机 <-> 激光转换测试
        laser_from_cam = camera_to_laser(camera_point)
        cam_from_laser = laser_to_camera(laser_point)
        
        print(f"  相机 -> 激光: {laser_from_cam}")
        print(f"  激光 -> 相机: {cam_from_laser}")
        print("")


def example_usage():
    """使用示例：展示如何使用坐标系转换函数"""
    print("使用示例：")
    print("注意：现在已经包含轴转换：")
    print("- 腰部→相机/激光：矩阵转换 → 逆向轴转换")
    print("- 相机/激光→腰部：轴转换 → 矩阵转换")
    print("- 激光雷达：绕x轴反转180°，y和z轴反向")
    print("- 深度相机：x→z, y→-x, z→-y")
    print()
    
    # 定义一个腰部坐标系中的点
    waist_point = [50.0, 30.0, 100.0]
    print(f"原始腰部坐标: {waist_point}")
    
    # 转换到相机坐标系（包含轴转换）
    camera_point = waist_to_camera(waist_point)
    print(f"转换到相机坐标（已轴转换）: {camera_point}")
    
    # 转换到激光坐标系（包含轴转换）
    laser_point = waist_to_laser(waist_point)
    print(f"转换到激光坐标（已轴转换）: {laser_point}")
    
    # 相机和激光之间的直接转换
    laser_from_camera = camera_to_laser(camera_point)
    print(f"从相机转换到激光: {laser_from_camera}")
    
    # 验证转换的一致性
    print(f"直接转换vs经腰部转换的误差: {np.linalg.norm(np.array(laser_point) - np.array(laser_from_camera)):.2e}")
    print()


if __name__ == "__main__":
    test_axis_transform()
    test_transform()
    print("="*50)
    example_usage()
